#!/usr/bin/env python3
import argparse
import os
import sys
from time import sleep, strftime
import json
import grpc

sys.path.append(
    os.path.join(os.path.dirname(os.path.abspath(__file__)),
                 '../../utils/'))
import p4runtime_lib.bmv2
import p4runtime_lib.helper
from p4runtime_lib.switch import ShutdownAllSwitchConnections

# 从拓扑文件加载配置
def load_topology(topology_file):
    with open(topology_file) as f:
        return json.load(f)

# 根据拓扑构建端口映射
def build_port_mapping(topology):
    port_map = {}
    for link in topology['links']:
        node1, node2 = link
        if 'h' in node1:  # 主机到交换机连接
            sw = node2.split('-')[0]
            port = int(node2.split('-p')[1])
            port_map[(sw, node1)] = port
        elif 'h' in node2:
            sw = node1.split('-')[0]
            port = int(node1.split('-p')[1])
            port_map[(sw, node2)] = port
        else:  # 交换机间连接
            sw1, port1 = node1.split('-p')
            sw2, port2 = node2.split('-p')
            port_map[(sw1, sw2)] = int(port1)
            port_map[(sw2, sw1)] = int(port2)
    return port_map

# 计数器日志记录
class CounterLogger:
    def __init__(self, log_dir="counter_logs"):
        self.log_dir = log_dir
        os.makedirs(log_dir, exist_ok=True)
        
    def log_counter(self, switch_name, counter_name, tunnel_id, packet_count, byte_count):
        timestamp = strftime("%Y-%m-%d %H:%M:%S")
        log_entry = f"{timestamp},{switch_name},{counter_name},{tunnel_id},{packet_count},{byte_count}\n"
        
        log_file = os.path.join(self.log_dir, f"{switch_name}_{counter_name}.csv")
        header = "timestamp,switch_name,counter_name,tunnel_id,packet_count,byte_count\n"
        
        if not os.path.exists(log_file):
            with open(log_file, 'w') as f:
                f.write(header)
        
        with open(log_file, 'a') as f:
            f.write(log_entry)

def readTableRules(p4info_helper, sw):
    """读取并打印交换机上的所有表项"""
    print(f'\n----- Reading table rules for {sw.name} -----')
    for response in sw.ReadTableEntries():
        for entity in response.entities:
            entry = entity.table_entry
            table_name = p4info_helper.get_tables_name(entry.table_id)
            print(f'{table_name}: ', end=' ')
            for m in entry.match:
                print(f"{p4info_helper.get_match_field_name(table_name, m.field_id)}="
                      f"{p4info_helper.get_match_field_value(m)}", end=' ')
            action = entry.action.action
            action_name = p4info_helper.get_actions_name(action.action_id)
            print(f'-> {action_name}', end=' ')
            for p in action.params:
                print(f"{p4info_helper.get_action_param_name(action_name, p.param_id)}="
                      f"{p.value}", end=' ')
            print()

def writeTunnelRules(p4info_helper, ingress_sw, egress_sw, tunnel_id,
                     dst_eth_addr, dst_ip_addr, port_map, host_mapping):
    """
    安装隧道规则，根据拓扑动态确定端口
    """
    # 1) 入口规则 (封装)
    table_entry = p4info_helper.buildTableEntry(
        table_name="MyIngress.ipv4_lpm",
        match_fields={"hdr.ipv4.dstAddr": (dst_ip_addr, 32)},
        action_name="MyIngress.myTunnel_ingress",
        action_params={"dst_id": tunnel_id})
    ingress_sw.WriteTableEntry(table_entry)
    print(f"Installed ingress tunnel rule on {ingress_sw.name}")

    # 2) 中转规则 (转发)
    next_hop = egress_sw.name
    forward_port = port_map.get((ingress_sw.name, next_hop))
    if forward_port is None:
        print(f"Warning: No port mapping found from {ingress_sw.name} to {next_hop}")
        return

    table_entry = p4info_helper.buildTableEntry(
        table_name="MyIngress.myTunnel_exact",
        match_fields={"hdr.myTunnel.dst_id": tunnel_id},
        action_name="MyIngress.myTunnel_forward",
        action_params={"port": forward_port})
    ingress_sw.WriteTableEntry(table_entry)
    print(f"Installed transit tunnel rule on {ingress_sw.name} (port {forward_port})")

    # 3) 出口规则 (解封装)
    # 使用host_mapping来获取正确的主机名
    host_name = host_mapping.get(tunnel_id)
    if host_name is None:
        print(f"Warning: No host mapping found for tunnel ID {tunnel_id}")
        return

    egress_port = port_map.get((egress_sw.name, host_name))
    if egress_port is None:
        print(f"Warning: No port mapping found from {egress_sw.name} to {host_name}")
        return

    table_entry = p4info_helper.buildTableEntry(
        table_name="MyIngress.myTunnel_exact",
        match_fields={"hdr.myTunnel.dst_id": tunnel_id},
        action_name="MyIngress.myTunnel_egress",
        action_params={
            "dstAddr": dst_eth_addr,
            "port": egress_port})
    egress_sw.WriteTableEntry(table_entry)
    print(f"Installed egress tunnel rule on {egress_sw.name} (port {egress_port})")

def setup_all_tunnels(p4info_helper, switches, topology, port_map):
    """配置所有主机间的双向隧道"""
    hosts = topology['hosts']
    tunnel_rules = []
    
    # 隧道ID到主机名的映射
    host_mapping = {
        100: 'h2', 101: 'h1',  # h1 <-> h2
        200: 'h3', 201: 'h1',  # h1 <-> h3
        300: 'h3', 301: 'h2'   # h2 <-> h3
    }
    
    # h1 <-> h2 (tunnel 100/101)
    tunnel_rules.append({
        'ingress': 's1', 'egress': 's2', 'id': 100,
        'eth': hosts['h2']['mac'], 'ip': hosts['h2']['ip'].split('/')[0]
    })
    tunnel_rules.append({
        'ingress': 's2', 'egress': 's1', 'id': 101,
        'eth': hosts['h1']['mac'], 'ip': hosts['h1']['ip'].split('/')[0]
    })
    
    # h1 <-> h3 (tunnel 200/201)
    tunnel_rules.append({
        'ingress': 's1', 'egress': 's3', 'id': 200,
        'eth': hosts['h3']['mac'], 'ip': hosts['h3']['ip'].split('/')[0]
    })
    tunnel_rules.append({
        'ingress': 's3', 'egress': 's1', 'id': 201,
        'eth': hosts['h1']['mac'], 'ip': hosts['h1']['ip'].split('/')[0]
    })
    
    # h2 <-> h3 (tunnel 300/301)
    tunnel_rules.append({
        'ingress': 's2', 'egress': 's3', 'id': 300,
        'eth': hosts['h3']['mac'], 'ip': hosts['h3']['ip'].split('/')[0]
    })
    tunnel_rules.append({
        'ingress': 's3', 'egress': 's2', 'id': 301,
        'eth': hosts['h2']['mac'], 'ip': hosts['h2']['ip'].split('/')[0]
    })

    for rule in tunnel_rules:
        writeTunnelRules(
            p4info_helper,
            ingress_sw=switches[rule['ingress']],
            egress_sw=switches[rule['egress']],
            tunnel_id=rule['id'],
            dst_eth_addr=rule['eth'],
            dst_ip_addr=rule['ip'],
            port_map=port_map,
            host_mapping=host_mapping
        )
    
    return tunnel_rules

def monitor_counters(p4info_helper, switches, tunnel_rules, counter_logger):
    """监控并记录计数器"""
    try:
        while True:
            sleep(2)
            print('\n----- Reading tunnel counters -----')
            
            for rule in tunnel_rules:
                ingress_sw = switches[rule['ingress']]
                egress_sw = switches[rule['egress']]
                tunnel_id = rule['id']
                
                # 入口计数器
                for response in ingress_sw.ReadCounters(
                    p4info_helper.get_counters_id("MyIngress.ingressTunnelCounter"), 
                    tunnel_id):
                    for entity in response.entities:
                        counter = entity.counter_entry
                        print(f"{ingress_sw.name} ingress {tunnel_id}: "
                              f"{counter.data.packet_count} packets")
                        counter_logger.log_counter(
                            ingress_sw.name, "ingress", tunnel_id,
                            counter.data.packet_count, counter.data.byte_count)
                
                # 出口计数器
                for response in egress_sw.ReadCounters(
                    p4info_helper.get_counters_id("MyIngress.egressTunnelCounter"), 
                    tunnel_id):
                    for entity in response.entities:
                        counter = entity.counter_entry
                        print(f"{egress_sw.name} egress {tunnel_id}: "
                              f"{counter.data.packet_count} packets")
                        counter_logger.log_counter(
                            egress_sw.name, "egress", tunnel_id,
                            counter.data.packet_count, counter.data.byte_count)
    
    except KeyboardInterrupt:
        print("\nStopping counter monitoring")

def main(p4info_file_path, bmv2_file_path, topology_file):
    # 加载拓扑配置
    topology = load_topology(topology_file)
    port_map = build_port_mapping(topology)
    
    # 初始化P4Runtime帮助类
    p4info_helper = p4runtime_lib.helper.P4InfoHelper(p4info_file_path)
    counter_logger = CounterLogger()

    try:
        # 创建交换机连接
        switches = {
            's1': p4runtime_lib.bmv2.Bmv2SwitchConnection(
                name='s1', address='127.0.0.1:50051', device_id=0,
                proto_dump_file='logs/s1-p4runtime-requests.txt'),
            's2': p4runtime_lib.bmv2.Bmv2SwitchConnection(
                name='s2', address='127.0.0.1:50052', device_id=1,
                proto_dump_file='logs/s2-p4runtime-requests.txt'),
            's3': p4runtime_lib.bmv2.Bmv2SwitchConnection(
                name='s3', address='127.0.0.1:50053', device_id=2,
                proto_dump_file='logs/s3-p4runtime-requests.txt')
        }

        # 建立主仲裁连接
        for sw in switches.values():
            sw.MasterArbitrationUpdate()
            sw.SetForwardingPipelineConfig(
                p4info=p4info_helper.p4info,
                bmv2_json_file_path=bmv2_file_path)
            print(f"Installed P4 program on {sw.name}")

        # 设置所有隧道规则
        tunnel_rules = setup_all_tunnels(p4info_helper, switches, topology, port_map)

        # 读取并显示表项 - 使用我们自己的函数
        for sw in switches.values():
            readTableRules(p4info_helper, sw)

        # 开始监控计数器
        monitor_counters(p4info_helper, switches, tunnel_rules, counter_logger)

    except grpc.RpcError as e:
        print(f"gRPC Error: {e.details()}")
    finally:
        ShutdownAllSwitchConnections()

if __name__ == '__main__':
    parser = argparse.ArgumentParser(description='P4Runtime Controller')
    parser.add_argument('--p4info', help='p4info proto in text format from p4c',
                       default='./build/advanced_tunnel.p4.p4info.txt')
    parser.add_argument('--bmv2-json', help='BMv2 JSON file from p4c',
                       default='./build/advanced_tunnel.json')
    parser.add_argument('--topo', help='Topology JSON file',
                       default='topology.json')
    args = parser.parse_args()

    if not all(os.path.exists(f) for f in [args.p4info, args.bmv2_json, args.topo]):
        print("Error: Required files not found. Please check paths.")
        sys.exit(1)

    main(args.p4info, args.bmv2_json, args.topo)